package cn.ichiva.dl4j.gnn.common;

import cn.ichiva.dl4j.gnn.mnist.MnistAgent;
import org.nd4j.linalg.api.ndarray.INDArray;

import java.io.*;
import java.util.HashMap;
import java.util.Map;

public class Warehouse {

    private static final File dir;
    static {
        dir = new File("./dna");
        if (!dir.exists()) {
            dir.mkdirs();
        }
    }

    public static void write(MnistAgent agent) throws IOException {
        Map<String,Object> map = new HashMap<>();
        map.put("score",agent.getScore());
        map.put("validScore",agent.getValidScore());
        map.put("dna",agent.getDna());

        File file = new File(dir,"tthd-" + (int)(agent.getScore() * 10000) +".objs");
        try (ObjectOutputStream out = new ObjectOutputStream(new FileOutputStream(file))){
            out.writeObject(map);
            out.flush();
        }
    }

    public static String[] readList(){
        return dir.list();
    }

    public static MnistAgent readAgent(String name) throws IOException, ClassNotFoundException {
        File file = new File(dir,name);
        try (ObjectInputStream in = new ObjectInputStream(new FileInputStream(file))){
            Map<String,Object> map = (Map<String, Object>) in.readObject();
            MnistAgent agent = MnistAgent.create((INDArray) map.get("dna"));
            agent.setScore((Double) map.get("score"));
            agent.setValidScore((Double) map.get("validScore"));
            return agent;
        }
    }

}
